import dnnlib
import legacy
import numpy as np
import torch
from PIL import Image


def generate_images(network_pkl, num_images=1, truncation_psi=0.7):
    """
    使用 StyleGAN2 生成图像
    
    参数:
        network_pkl: 预训练模型路径
        num_images: 要生成的图像数量
        truncation_psi: truncation 参数，控制生成图像的变化程度
    """
    
    # 加载预训练模型
    print(f'加载网络: {network_pkl}')
    with dnnlib.util.open_url(network_pkl) as f:
        G = legacy.load_network_pkl(f)['G_ema'].cuda()

    # 生成随机潜码
    z = torch.randn([num_images, G.z_dim]).cuda()
    
    # 生成图像
    print('生成图像...')
    with torch.no_grad():
        imgs = G(z, None, truncation_psi=truncation_psi)
        imgs = (imgs.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).cpu().numpy()

    # 保存图像
    for i, img in enumerate(imgs):
        PIL.Image.fromarray(img, 'RGB').save(f'generated_{i:04d}.png')

    print('完成！')

if __name__ == '__main__':
    # 使用 FFHQ 预训练模型示例
    network_pkl = "https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/ffhq.pkl"
    
    # 生成 3 张图像
    generate_images(network_pkl, num_images=3)